feat: add out-of-tree and mcp based classification support #375
Conversation
Signed-off-by: Huamin Chen <[email protected]>
✅ Deploy Preview for vllm-semantic-router ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
👥 vLLM Semantic Team NotificationThe following members have been identified for the changed files in this PR and have been automatically assigned: 📁
|
|
Will add a simple mcp classification server for testing next. |
There was a problem hiding this comment.
Pull Request Overview
This PR introduces MCP (Model Context Protocol) based classification support as an alternative to in-tree classification models. It enables external classification services for category determination while maintaining the existing in-tree classifier as a fallback option.
- Adds comprehensive MCP client implementation with stdio and HTTP transport support
- Integrates MCP-based category classifier into the existing classification pipeline
- Provides flexible configuration options for choosing between in-tree and MCP classification
Reviewed Changes
Copilot reviewed 11 out of 12 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| src/semantic-router/pkg/utils/classification/mcp_classifier.go | Implements MCP-based category classification with full probability distribution support |
| src/semantic-router/pkg/utils/classification/classifier.go | Updates classifier to support both in-tree and MCP classification methods |
| src/semantic-router/pkg/connectivity/mcp/*.go | Complete MCP client implementation with factory pattern and multiple transport types |
| src/semantic-router/pkg/config/config.go | Adds configuration structure for MCP category classifier |
| src/semantic-router/go.mod | Adds mcp-go dependency for MCP protocol support |
| config/config-mcp-classifier-example.yaml | Provides example configuration for MCP-based classification |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| reasoningDecision.Confidence, | ||
| reasoningDecision.UseReasoning, | ||
| reasoningDecision.DecisionReason, | ||
| topCategory, |
There was a problem hiding this comment.
The topCategory variable is being passed twice as arguments (lines 402 and 408) to RecordEntropyClassificationMetrics. This appears to be a copy-paste error where the second occurrence should likely be a different parameter.
Copilot uses AI. Check for mistakes.
| // NOTE: Commented out - promptGuard is in BaseClient which is commented out | ||
| // if c.promptGuard != nil && c.promptGuard.IsEnabled() { | ||
| if false { | ||
| if err := (error)(nil); err != nil { // c.promptGuard.ValidateToolCall(name) |
There was a problem hiding this comment.
This disabled code block with hardcoded false condition and dummy error assignment should be removed or properly implemented. The commented-out functionality suggests incomplete security validation that could be important for production use.
| // NOTE: Commented out - promptGuard is in BaseClient which is commented out | |
| // if c.promptGuard != nil && c.promptGuard.IsEnabled() { | |
| if false { | |
| if err := (error)(nil); err != nil { // c.promptGuard.ValidateToolCall(name) | |
| // Validate tool call using promptGuard if enabled | |
| if c.promptGuard != nil && c.promptGuard.IsEnabled() { | |
| if err := c.promptGuard.ValidateToolCall(name); err != nil { |
Copilot uses AI. Check for mistakes.
|
|
||
| // Try to initialize the connection | ||
| initRequest := mcp.InitializeRequest{} | ||
| initRequest.Params.ProtocolVersion = "2024-11-05" |
There was a problem hiding this comment.
The hardcoded protocol version should be defined as a constant to ensure consistency across the codebase and make version updates easier to manage.
Copilot uses AI. Check for mistakes.
| func ImageContentBlock(data, mimeType string) ContentBlock { | ||
| return ContentBlock{ | ||
| Type: "image", | ||
| Text: data, // In practice, this would be base64 encoded data |
There was a problem hiding this comment.
The comment indicates this field should contain base64 encoded data, but the function parameter is named data without specifying the expected format. This could lead to incorrect usage where non-base64 data is passed.
| func ImageContentBlock(data, mimeType string) ContentBlock { | |
| return ContentBlock{ | |
| Type: "image", | |
| Text: data, // In practice, this would be base64 encoded data | |
| func ImageContentBlock(base64Data, mimeType string) ContentBlock { | |
| return ContentBlock{ | |
| Type: "image", | |
| Text: base64Data, // base64-encoded image data |
Copilot uses AI. Check for mistakes.
| // Legacy types moved to interface.go - keeping for backward compatibility | ||
|
|
||
| // NewLegacyClient creates a legacy MCP client wrapper (deprecated - use factory instead) | ||
| func NewLegacyClient(name string, config ClientConfig) *Client { |
There was a problem hiding this comment.
The function is marked as deprecated but still implemented. Consider adding a deprecation warning log or removing this function entirely if the factory pattern is the preferred approach.
| func NewLegacyClient(name string, config ClientConfig) *Client { | |
| func NewLegacyClient(name string, config ClientConfig) *Client { | |
| log.Printf("[DEPRECATION WARNING] NewLegacyClient is deprecated and will be removed in a future release. Please use the factory pattern instead.") |
Copilot uses AI. Check for mistakes.
| } | ||
|
|
||
| // Add MCP classifier if configured (only if in-tree is not configured) | ||
| if cfg.Classifier.MCPCategoryModel.Enabled && cfg.Classifier.CategoryModel.ModelID == "" { |
There was a problem hiding this comment.
The condition that MCP classifier is only enabled when in-tree classifier is disabled (ModelID is empty) creates a mutually exclusive relationship that could be confusing. Consider making this relationship more explicit or allowing both to coexist with clear priority rules.
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| return authHeader | ||
| } | ||
|
|
||
| // SSEMessage represents a Server-Sent Events message |
There was a problem hiding this comment.
Incomplete documentation - the SSEMessage type is declared in the comment but the actual struct definition is missing from the file.
| // SSEMessage represents a Server-Sent Events message | |
| // SSEMessage represents a Server-Sent Events message | |
| type SSEMessage struct { | |
| Event string `json:"event,omitempty"` | |
| Data interface{} `json:"data"` | |
| ID string `json:"id,omitempty"` | |
| Retry int `json:"retry,omitempty"` | |
| } |
Copilot uses AI. Check for mistakes.
| // Start the client (required before Initialize) | ||
| // NOTE: Start() method removed in mcp-go v0.9.0 - connection starts automatically | ||
| // if err := c.mcpClient.Start(ctx); err != nil { | ||
| // c.log(LoggingLevelError, fmt.Sprintf("Failed to start MCP client: %v", err)) | ||
| // initDone <- fmt.Errorf("failed to start MCP client: %w", err) | ||
| // return | ||
| // } |
There was a problem hiding this comment.
Large block of commented-out code should be removed. If this is important for version compatibility documentation, consider moving it to a separate comment explaining the API change.
| // Start the client (required before Initialize) | |
| // NOTE: Start() method removed in mcp-go v0.9.0 - connection starts automatically | |
| // if err := c.mcpClient.Start(ctx); err != nil { | |
| // c.log(LoggingLevelError, fmt.Sprintf("Failed to start MCP client: %v", err)) | |
| // initDone <- fmt.Errorf("failed to start MCP client: %w", err) | |
| // return | |
| // } | |
| // Start() method removed in mcp-go v0.9.0 - connection starts automatically |
Copilot uses AI. Check for mistakes.
| initRequest.Params.ClientInfo = mcp.Implementation{ | ||
| Name: "http-mcp-client", | ||
| Version: "1.0.0", | ||
| } |
There was a problem hiding this comment.
Hardcoded version '1.0.0' should be extracted to a constant or configuration parameter to maintain consistency and enable easier version management.
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated 4 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| // Set timeout if specified | ||
| if cfg.Classifier.MCPCategoryModel.TimeoutSeconds > 0 { |
There was a problem hiding this comment.
Consider adding validation to ensure TimeoutSeconds is within reasonable bounds (e.g., greater than 0 and less than a maximum threshold) to prevent indefinite timeouts or zero-duration timeouts.
| // Set timeout if specified | |
| if cfg.Classifier.MCPCategoryModel.TimeoutSeconds > 0 { | |
| // Set timeout if specified and within reasonable bounds | |
| const maxTimeoutSeconds = 300 // 5 minutes | |
| if cfg.Classifier.MCPCategoryModel.TimeoutSeconds > 0 { | |
| if cfg.Classifier.MCPCategoryModel.TimeoutSeconds > maxTimeoutSeconds { | |
| return fmt.Errorf("TimeoutSeconds (%d) exceeds maximum allowed (%d seconds)", cfg.Classifier.MCPCategoryModel.TimeoutSeconds, maxTimeoutSeconds) | |
| } |
Copilot uses AI. Check for mistakes.
| // Check threshold | ||
| threshold := c.Config.Classifier.MCPCategoryModel.Threshold | ||
| if threshold == 0 { | ||
| threshold = 0.5 // Default threshold |
There was a problem hiding this comment.
The default threshold value (0.5) is hardcoded in multiple places in the file (lines 264 and 345). Consider defining this as a constant to ensure consistency and easier maintenance.
Copilot uses AI. Check for mistakes.
| }() | ||
|
|
||
| c.log(LoggingLevelDebug, fmt.Sprintf("Starting MCP client...")) | ||
| // Note: Start() method removed in mcp-go v0.9.0+ - connection starts automatically |
There was a problem hiding this comment.
This comment indicates the code is compatible with mcp-go v0.9.0+, but the go.mod file shows v0.42.0-beta.1. The comment should be updated to reflect the actual version being used or clarify the version compatibility.
| // Note: Start() method removed in mcp-go v0.9.0+ - connection starts automatically | |
| // Note: In mcp-go v0.42.0-beta.1 (current), connection starts automatically; Start() method is not used. |
Copilot uses AI. Check for mistakes.
| // The mcp.Content type doesn't have a Type field, so we need to use type assertion directly | ||
| firstContent := result.Content[0] | ||
|
|
||
| // Try to extract content based on the actual type | ||
| if textContent, ok := firstContent.(*mcp.TextContent); ok { | ||
| content = textContent.Text | ||
| } else if imageContent, ok := firstContent.(*mcp.ImageContent); ok { | ||
| content = fmt.Sprintf("Image: %s", imageContent.Data) | ||
| } else if resourceContent, ok := firstContent.(*mcp.EmbeddedResource); ok { | ||
| // The mcp.EmbeddedResource has Resource field but it's an interface, so we'll use string representation | ||
| content = fmt.Sprintf("Resource: %v", resourceContent.Resource) | ||
| } else { | ||
| // Fallback: try to get string representation |
There was a problem hiding this comment.
The type assertions for mcp.TextContent, mcp.ImageContent, and mcp.EmbeddedResource may not work correctly with the actual mcp library types. These types should be verified against the actual mcp-go library API to ensure proper type matching.
| // The mcp.Content type doesn't have a Type field, so we need to use type assertion directly | |
| firstContent := result.Content[0] | |
| // Try to extract content based on the actual type | |
| if textContent, ok := firstContent.(*mcp.TextContent); ok { | |
| content = textContent.Text | |
| } else if imageContent, ok := firstContent.(*mcp.ImageContent); ok { | |
| content = fmt.Sprintf("Image: %s", imageContent.Data) | |
| } else if resourceContent, ok := firstContent.(*mcp.EmbeddedResource); ok { | |
| // The mcp.EmbeddedResource has Resource field but it's an interface, so we'll use string representation | |
| content = fmt.Sprintf("Resource: %v", resourceContent.Resource) | |
| } else { | |
| // Fallback: try to get string representation | |
| firstContent := result.Content[0] | |
| // Use a type switch to match the actual types from mcp-go | |
| switch c := firstContent.(type) { | |
| case *mcp.TextContent: | |
| content = c.Text | |
| case mcp.TextContent: | |
| content = c.Text | |
| case *mcp.ImageContent: | |
| content = fmt.Sprintf("Image: %s", c.Data) | |
| case mcp.ImageContent: | |
| content = fmt.Sprintf("Image: %s", c.Data) | |
| case *mcp.EmbeddedResource: | |
| content = fmt.Sprintf("Resource: %v", c.Resource) | |
| case mcp.EmbeddedResource: | |
| content = fmt.Sprintf("Resource: %v", c.Resource) | |
| default: |
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| var arguments map[string]interface{} | ||
| if openAICall.Function.Arguments != "" { | ||
| if err := json.Unmarshal([]byte(openAICall.Function.Arguments), &arguments); err != nil { | ||
| return mcp.CallToolRequest{}, fmt.Errorf("failed to parse arguments: %w", err) |
There was a problem hiding this comment.
The error message for JSON parsing should provide more context about what failed to parse.
| return mcp.CallToolRequest{}, fmt.Errorf("failed to parse arguments: %w", err) | |
| argStr := openAICall.Function.Arguments | |
| const maxLen = 200 | |
| if len(argStr) > maxLen { | |
| argStr = argStr[:maxLen] + "...(truncated)" | |
| } | |
| return mcp.CallToolRequest{}, fmt.Errorf("failed to parse arguments (%q): %w", argStr, err) |
Copilot uses AI. Check for mistakes.
| ) | ||
|
|
||
| const ( | ||
| // DefaultMCPThreshold is the default confidence threshold for MCP classification |
There was a problem hiding this comment.
The DefaultMCPThreshold constant should have a comment explaining why 0.5 was chosen as the default value.
| // DefaultMCPThreshold is the default confidence threshold for MCP classification | |
| // DefaultMCPThreshold is the default confidence threshold for MCP classification. | |
| // A value of 0.5 is commonly used as it represents a neutral cutoff for binary or probabilistic classification, | |
| // meaning predictions with confidence above 50% are considered positive. Adjust as needed for your use case. |
Copilot uses AI. Check for mistakes.
| // Note: Both in-tree and MCP classifiers can be configured simultaneously. | ||
| // At runtime, in-tree classifier will be tried first, with MCP as a fallback. | ||
| // This allows flexible deployment scenarios (e.g., gradual migration, A/B testing). |
There was a problem hiding this comment.
This important fallback behavior should also be documented in the function-level documentation, not just as inline comments.
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| // A value of 0.5 is commonly used as it represents a neutral cutoff for binary or probabilistic classification, | ||
| // meaning predictions with confidence above 50% are considered positive. Adjust as needed for your use case. |
There was a problem hiding this comment.
The comment refers to 'binary or probabilistic classification' but this appears to be for multi-class category classification. The explanation about '50% positive' is misleading for multi-class scenarios where 0.5 threshold means the predicted class needs >50% confidence.
| // A value of 0.5 is commonly used as it represents a neutral cutoff for binary or probabilistic classification, | |
| // meaning predictions with confidence above 50% are considered positive. Adjust as needed for your use case. | |
| // For multi-class classification, a value of 0.5 means that a predicted class must have more than 50% confidence | |
| // to be selected. Adjust this threshold as needed for your use case. |
Copilot uses AI. Check for mistakes.
| c.Config.Classifier.MCPCategoryModel.ToolName != "" && | ||
| c.mcpCategoryInitializer != nil |
There was a problem hiding this comment.
The method checks if mcpCategoryInitializer is not nil, but this creates a circular dependency where the initializer must exist to determine if MCP is enabled. Consider separating configuration validation from runtime state checking.
| c.Config.Classifier.MCPCategoryModel.ToolName != "" && | |
| c.mcpCategoryInitializer != nil | |
| c.Config.Classifier.MCPCategoryModel.ToolName != "" |
Copilot uses AI. Check for mistakes.
| // MCPProtocolVersion is the MCP protocol version supported by this implementation | ||
| MCPProtocolVersion = "2024-11-05" |
There was a problem hiding this comment.
The hardcoded protocol version '2024-11-05' should be configurable or derived from the mcp-go library to ensure compatibility and easier updates.
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
|
|
||
| const ( | ||
| // DefaultMCPThreshold is the default confidence threshold for MCP classification. | ||
| // For multi-class classification, a value of 0.5 means that a predicted class must have more than 50% confidence |
There was a problem hiding this comment.
The documentation incorrectly states 'more than 50% confidence' when the code uses >= comparison (result.Confidence < threshold). A threshold of 0.5 means the confidence must be 0.5 or higher, not more than 50%.
| // For multi-class classification, a value of 0.5 means that a predicted class must have more than 50% confidence | |
| // For multi-class classification, a value of 0.5 means that a predicted class must have 50% or higher confidence |
Copilot uses AI. Check for mistakes.
| if result.Confidence < threshold { | ||
| observability.Infof("MCP classification confidence (%.4f) below threshold (%.4f)", | ||
| result.Confidence, threshold) | ||
| return "", float64(result.Confidence), nil | ||
| } |
There was a problem hiding this comment.
The threshold comparison logic conflicts with the comment at line 22. The code correctly implements 'confidence must be >= threshold' but the comment suggests 'more than threshold'.
Copilot uses AI. Check for mistakes.
| // NewClassifier creates a new classifier with model selection and jailbreak/PII detection capabilities. | ||
| // Both in-tree and MCP classifiers can be configured simultaneously for category classification. | ||
| // At runtime, in-tree classifier will be tried first, with MCP as a fallback, | ||
| // allowing flexible deployment scenarios such as gradual migration or A/B testing. |
There was a problem hiding this comment.
The documentation mentions 'A/B testing' as a use case, but the current implementation doesn't support A/B testing - it only provides a fallback mechanism where in-tree is always tried first. A/B testing would require random selection between classifiers.
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 17 out of 18 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
src/semantic-router/pkg/connectivity/mcp/types.go:1
- The comment claims mcp-go doesn't export a version constant, but this should be verified. Consider using a constant from the mcp-go library if available to maintain consistency and avoid version mismatches.
package mcp
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| }() | ||
|
|
||
| c.log(LoggingLevelDebug, fmt.Sprintf("Starting MCP client...")) | ||
| // Note: In mcp-go v0.42.0-beta.1, connection starts automatically; Start() method is not used |
There was a problem hiding this comment.
This version-specific comment may become outdated. Consider making this comment more generic or removing the specific version reference to improve maintainability.
| // Note: In mcp-go v0.42.0-beta.1, connection starts automatically; Start() method is not used | |
| // Note: Connection starts automatically; Start() method is not used |
Copilot uses AI. Check for mistakes.
| "--http", action="store_true", help="Run in HTTP mode instead of stdio" | ||
| ) | ||
| parser.add_argument( | ||
| "--port", type=int, default=8090, help="HTTP port (default: 8090)" |
There was a problem hiding this comment.
[nitpick] The help text says "default: 8090" but this information is redundant since argparse automatically shows the default value. Consider simplifying to just "HTTP port".
| "--port", type=int, default=8090, help="HTTP port (default: 8090)" | |
| "--port", type=int, default=8090, help="HTTP port" |
Copilot uses AI. Check for mistakes.
Signed-off-by: Huamin Chen <[email protected]>
| transport_type: "http" # HTTP transport | ||
| url: "http://localhost:8090/mcp" # MCP server endpoint | ||
|
|
||
| tool_name: "classify_text" # MCP tool name to call |
There was a problem hiding this comment.
why should we specify the tool_name in client?
There was a problem hiding this comment.
good point! the tool name is replaced by auto tool discovery now.
| } | ||
|
|
||
| // Parse JSON response: {"class": int, "confidence": float, "model": str, "use_reasoning": bool} | ||
| var response struct { |
There was a problem hiding this comment.
so i think we rely on the mcp server to respect the format right?
There was a problem hiding this comment.
should we move this into a standable api pkg? then the mcp server can export the pkg to follow the API.
There was a problem hiding this comment.
sure, this one is moved to types.go
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
Signed-off-by: Huamin Chen <[email protected]>
|
@Xunzhuo @wangchen615 PTAL. I am going to add the detailed MCP integration to the doc using the example server to explain how to add external classification support. |
3 participants
What type of PR is this?
Fix #368
MCP Classification Test
Run MCP Server
cd examples/mcp-classifier-server python server.py --http --port 8090Run Router
Test prompt
MCP Server log
Router log
What this PR does / why we need it:
Which issue(s) this PR fixes:
Fixes #
Release Notes: Yes/No